;;;
;;; ysmoke.scm - Real-Time Fluid Dynamics for Games for Ypsilon
;;;
;;;   Copyright (c) 2009 Teruaki Gemma <teruakigemma@gmail.com>
;;;
;;;   Redistribution and use in source and binary forms, with or without
;;;   modification, are permitted provided that the following conditions
;;;   are met:
;;;
;;;   1. Redistributions of source code must retain the above copyright
;;;      notice, this list of conditions and the following disclaimer.
;;;
;;;   2. Redistributions in binary form must reproduce the above copyright
;;;      notice, this list of conditions and the following disclaimer in the
;;;      documentation and/or other materials provided with the distribution.
;;;
;;;   3. Neither the name of the authors nor the names of its contributors
;;;      may be used to endorse or promote products derived from this
;;;      software without specific prior written permission.
;;;
;;;   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
;;;   "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
;;;   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
;;;   A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
;;;   OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
;;;   SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
;;;   TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
;;;   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
;;;   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
;;;   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
;;;   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
;;;
;;;  $Id$
;;;

;;;
;;; Ported from the paper "Real-Time Fluid Dynamics for Games"
;;; by Jos Stam(jstam@aw.sgi.com) at GDC2003.
;;; You can get the paper and the original codes from
;;; http://www.dgp.toronto.edu/people/stam/reality/Research/pub.html
;;;

(import (rnrs)
        (ypsilon gl)
        (ypsilon glut))

(define win-id #f)
(define win-x 512)
(define win-y 512)
(define N 64)       ;; grid resolution
(define dt 0.1)     ;; time step
(define diff 0.0)   ;; diffusion rate of the density
(define visc 0.0)   ;; viscosity of the fluid
(define force 5.0)  ;; scales the mouse movement that generate a force
(define source 100) ;; amount of density that will be deposited
(define dvel #f)
(define u #f)
(define u-prev #f)
(define v #f)
(define v-prev #f) 
(define dens #f)
(define dens-prev #f)
(define omx 0)
(define omy 0)
(define mx 0)
(define my 0)
(define mouse-down (make-vector 3 #f))

;; for the FPS counting.
(define GLUT_ELAPSED_TIME 700)
(define *t0* 0)
(define *frames* 0)

(define (clamp x a b)
  (min (max x a) b))

(define-syntax swap
  (syntax-rules ()
    ((_ a b)
     (let ((c a))
       (set! a b)
       (set! b c)))))

(define-syntax for-each-cell
  (syntax-rules ()
    ((_ i j start end body ...)
     (do ((i start (+ i 1))) ((>= i end))
       (do ((j start (+ j 1))) ((>= j end))
         body ...)))))

(define (make-matrix dim fill)
  (let ((a (make-vector dim)))
    (do ((i 0 (+ i 1))) ((= i dim) a)
      (vector-set! a i (make-vector dim fill)))))

(define (matrix-set! mat i j v)
  (vector-set! (vector-ref mat i) j v))

(define (matrix-ref mat i j)
  (vector-ref (vector-ref mat i) j))

(define (clear-data)
  (set! u         (make-matrix (+ N 2) 0.0))
  (set! u-prev    (make-matrix (+ N 2) 0.0))
  (set! v         (make-matrix (+ N 2) 0.0))
  (set! v-prev    (make-matrix (+ N 2) 0.0))
  (set! dens      (make-matrix (+ N 2) 0.0))
  (set! dens-prev (make-matrix (+ N 2) 0.0)))

(define (add-source N x s dt)
  (for-each-cell i j 0 (+ N 2)
    (matrix-set! x i j (+ (matrix-ref x i j) (* dt (matrix-ref s i j))))))

(define (set-bnd N b x)
  (define N1 (+ N 1))
  (do ((i 1 (+ i 1))) ((> i N))
    (matrix-set! x 0  i (* (if (= b 1) -1 1) (matrix-ref x 1 i)))
    (matrix-set! x N1 i (* (if (= b 1) -1 1) (matrix-ref x N i)))
    (matrix-set! x i  0 (* (if (= b 2) -1 1) (matrix-ref x i 1)))
    (matrix-set! x i N1 (* (if (= b 2) -1 1) (matrix-ref x i N))))
  (matrix-set! x  0  0 (* 0.5 (+ (matrix-ref x 1  0) (matrix-ref x  0 1))))
  (matrix-set! x  0 N1 (* 0.5 (+ (matrix-ref x 1 N1) (matrix-ref x  0 N))))
  (matrix-set! x N1  0 (* 0.5 (+ (matrix-ref x N  0) (matrix-ref x N1 1))))
  (matrix-set! x N1 N1 (* 0.5 (+ (matrix-ref x N N1) (matrix-ref x N1 N)))))

(define (lin-solve N b x x0 a c)
  (do ((k 0 (+ k 1))) ((= k 20))
    (for-each-cell i j 1 (+ N 1)
      (matrix-set! x i j (/ (+ (matrix-ref x0 i j)
                               (* a (+ (matrix-ref x (- i 1) j)
                                       (matrix-ref x (+ i 1) j)
                                       (matrix-ref x i (- j 1))
                                       (matrix-ref x i (+ j 1)))))
                            c)))
    (set-bnd N b x)))

(define (diffuse N b x x0 diff dt)
  (define a (* dt diff N N))
  (lin-solve N b x x0 a (+ 1 (* 4 a))))

(define (advect N b d d0 u v dt)
  (define dt0 (* dt N))
  (for-each-cell i j 1 (+ N 1)
    (let* ((x (clamp (- i (* dt0 (matrix-ref u i j))) 0.5 (+ N 0.5)))
           (y (clamp (- j (* dt0 (matrix-ref v i j))) 0.5 (+ N 0.5)))
           (i0 (inexact->exact (floor x)))
           (i1 (+ i0 1))
           (j0 (inexact->exact (floor y)))
           (j1 (+ j0 1))
           (s1 (- x i0))
           (s0 (- 1 s1))
           (t1 (- y j0))
           (t0 (- 1 t1)))
      (matrix-set! d i j (+ (* s0 (+ (* t0 (matrix-ref d0 i0 j0))
                                     (* t1 (matrix-ref d0 i0 j1))))
                            (* s1 (+ (* t0 (matrix-ref d0 i1 j0))
                                     (* t1 (matrix-ref d0 i1 j1))))))))
  (set-bnd N b d))

(define (project N u v p div)
  (for-each-cell i j 1 (+ N 1)
    (matrix-set! div i j (/ (* -0.5 (+ (- (matrix-ref u (+ i 1) j)
                                          (matrix-ref u (- i 1) j))
                                       (- (matrix-ref v i (+ j 1))
                                          (matrix-ref v i (- j 1)))))
                            N))
    (matrix-set! p i j 0))
  (set-bnd N 0 div)
  (set-bnd N 0 p)
  (lin-solve N 0 p div 1 4)
  (for-each-cell i j 1 (+ N 1)
    (matrix-set! u i j (- (matrix-ref u i j)
                          (* 0.5 N (- (matrix-ref p (+ i 1) j)
                                      (matrix-ref p (- i 1) j)))))
    (matrix-set! v i j (- (matrix-ref v i j)
                          (* 0.5 N (- (matrix-ref p i (+ j 1))
                                      (matrix-ref p i (- j 1)))))))
  (set-bnd N 1 u)
  (set-bnd N 2 v))

(define (dens-step N x x0 u v diff dt)
  (add-source N x x0 dt)
  (swap x0 x)
  (diffuse N 0 x x0 diff dt)
  (swap x0 x)
  (advect N 0 x x0 u v dt))

(define (vel-step N u v u0 v0 visc dt)
  (add-source N u u0 dt)
  (add-source N v v0 dt)
  (swap u0 u)
  (diffuse N 1 u u0 visc dt)
  (swap v0 v)
  (diffuse N 2 v v0 visc dt)
  (project N u v u0 v0)
  (swap u0 u)
  (swap v0 v)
  (advect N 1 u u0 u0 v0 dt)
  (advect N 2 v v0 u0 v0 dt)
  (project N u v u0 v0))

(define (draw-velocity)
  (define h (/ 1.0 N))
  (glColor3f 1 1 1)
  (glLineWidth 1)
  (glBegin GL_LINES)
  (for-each-cell i j 1 (+ N 1)
     (let ((x (* h (- i 0.5)))
           (y (* h (- j 0.5))))
       (glVertex2f x y)
       (glVertex2f (+ x (matrix-ref u i j))
                   (+ y (matrix-ref v i j)))))
  (glEnd))

(define (draw-density)
  (define h (/ 1.0 N))
  (glBegin GL_QUADS)
  (for-each-cell i j 0 (+ N 1)
    (let ((x (* h (- i 0.5)))
          (y (* h (- j 0.5)))
          (d00 (matrix-ref dens i       j))
          (d01 (matrix-ref dens i       (+ j 1)))
          (d10 (matrix-ref dens (+ i 1) j))
          (d11 (matrix-ref dens (+ i 1) (+ j 1))))
      (glColor3f d00 d00 d00)
      (glVertex2f x       y)
      (glColor3f d10 d10 d10)
      (glVertex2f (+ x h) y)
      (glColor3f d11 d11 d11)
      (glVertex2f (+ x h) (+ y h))
      (glColor3f d01 d01 d01)
      (glVertex2f x       (+ y h))))
  (glEnd))


(define (pre-display)
  (glViewport 0 0 win-x win-y)
  (glMatrixMode GL_PROJECTION)
  (glLoadIdentity)
  (glOrtho 0 1 0 1 -1 1)
  (glClearColor 0 0 0 1)
  (glClear GL_COLOR_BUFFER_BIT))

(define (get-from-UI d u v)
  (for-each-cell i j 0 (+ N 2)
    (matrix-set! u i j 0.0)
    (matrix-set! v i j 0.0)
    (matrix-set! d i j 0.0))
  (when (or (vector-ref mouse-down 0)
            (vector-ref mouse-down 2))
    (let ((i (floor (* N (/ mx win-x))))
          (j (floor (* N (/ (- win-y my) win-y)))))
      (when (and (>= i 1) (<= i N) (>= j 1) (<= j N))
        (when (vector-ref mouse-down 0)
          (matrix-set! u i j (* force (- mx omx)))
          (matrix-set! v i j (* force (- omy my))))
        (when (vector-ref mouse-down 2)
          (matrix-set! d i j source))
        (set! omx mx)
        (set! omy my)))))

(define (key rk x y)
  (define (c-int->c-uchar c)
    (bitwise-and c #xff))
  (let ((k (c-int->c-uchar rk)))
    (cond
     ((= k (char->integer #\c)) (clear-data))
     ((= k (char->integer #\v)) (set! dvel (not dvel)))
     ((= k (char->integer #\q)) (exit)))))

(define (mouse button state x y)
  (set! omx x)
  (set! omy y)
  (set! mx x)
  (set! my y)
  (vector-set! mouse-down button (= state GLUT_DOWN)))

(define (motion x y)
  (set! mx x)
  (set! my y))

(define (reshape width height)
  (glutSetWindow win-id)
  (glutReshapeWindow width height)
  (set! win-x width)
  (set! win-y height))

(define (idle)
  (get-from-UI dens-prev u-prev v-prev)
  (vel-step N u v u-prev v-prev visc dt)
  (dens-step N dens dens-prev u v diff dt)
  (glutSetWindow win-id)
  (glutPostRedisplay))

(define (display)
  (pre-display)
  (if dvel (draw-velocity) (draw-density))
  (glutSwapBuffers)

    ;; for the FPS counting.
  (set! *frames* (+ 1 *frames*))
  (let ((t (glutGet GLUT_ELAPSED_TIME)))
    (when (>= (- t *t0*) 5000)
      (let ((seconds (/ (- t *t0*) 1000.0)))
        (format #t "~d in ~d seconds = ~d FPS~%" *frames* seconds (/ *frames* seconds))
        (set! *t0*     t)
        (set! *frames* 0)))))

(define (open-glut-window)
  (glutInitDisplayMode (bitwise-ior GLUT_DOUBLE GLUT_RGBA))

  (glutInitWindowPosition 0     0)
  (glutInitWindowSize     win-x win-y)

  (set! win-id (glutCreateWindow "Real-Time Fluid Dynamics for Games"))
  (glClearColor 0 0 0 1)
  (glClear GL_COLOR_BUFFER_BIT)
  (glutSwapBuffers)
  (glClear GL_COLOR_BUFFER_BIT)
  (glutSwapBuffers)

  (pre-display)

  (glutKeyboardFunc   key)
  (glutMouseFunc      mouse)
  (glutMotionFunc     motion)
  (glutReshapeFunc    reshape)
  (glutIdleFunc       idle)
  (glutDisplayFunc    display))

(begin
  (format #t "\n\nHow to use this demo:\n\n")
  (format #t "\t Add densities with the right mouse button\n" )
	(format #t "\t Add velocities with the left mouse button and dragging the mouse\n" )
	(format #t "\t Toggle density/velocity display with the 'v' key\n" )
	(format #t "\t Clear the simulation by pressing the 'c' key\n" )
	(format #t "\t Quit by pressing the 'q' key\n" )

  (glutInit (vector (length (command-line))) (apply vector (command-line)))
  (clear-data)
  (open-glut-window)
  (glutMainLoop)
  0)

